# -*- coding: utf-8 -*-
# write by hashaki
# first edit on 2018/11/01 last change on 2019/4/21
# all stander insicator

# fix by jack
# change by hasaki
"""
输入的data为list, 里面的元素为dict, key有open,close等
talib中文文档：https://github.com/HuaRongSAO/talib-document
"""
import talib
import numpy as np

# ---------------------------------------------------------
def transferData(data):
    """数据转化函数"""
    # 储存字典
    arrDict = dict()
    # OHLC转化
    arrDict["open"] = np.array([x["open"] for x in data])
    arrDict["high"] = np.array([x["high"] for x in data])
    arrDict["low"] = np.array([x["low"] for x in data])
    arrDict["close"] = np.array([x["close"] for x in data])
    arrDict["volume"] = np.array([x["volume"] for x in data])
    # 返回
    return arrDict


# ---------------------------------------------------------
def ma(data, n, array=False):
    """简单移动平均线"""
    arrDict = transferData(data) # 数据转化
    result = talib.SMA(arrDict["close"], timeperiod=n)
    if array:
        return result
    return result[-1]

# ---------------------------------------------------------
def ema(data, n, array=False):
    """指数移动平均线"""
    arrDict = transferData(data)  # 数据转化
    result = talib.EMA(arrDict["close"], timeperiod=n)
    if array:
        return result
    return result[-1]

# ---------------------------------------------------------
def kama(data, n, array=False):
    """考夫曼的自适应移动平均线"""
    arrDict = transferData(data)  # 数据转化
    result = talib.KAMA(arrDict["close"], timeperiod=n)
    if array:
        return result
    return result[-1]

# ---------------------------------------------------------
def wma(data, n, array=False):
    """加权移动平均线"""
    arrDict = transferData(data)  # 数据转化
    result = talib.WMA(arrDict["close"], timeperiod=n)
    if array:
        return result
    return result[-1]

# ---------------------------------------------------------
def sar(data, acceleration, maximum, array=False):
    """抛物线指标"""
    arrDict = transferData(data)  # 数据转化
    result = talib.SAR(arrDict["high"], arrDict["low"],
                       acceleration=acceleration, maximum=maximum)
    if array:
        return result
    return result[-1]

# ---------------------------------------------------------
def boll(data, n, dev, array=False):
    """布林通道, 上下对称"""
    arrDict = transferData(data)  # 数据转化
    upperband, middleband, lowerband = talib.BBANDS(arrDict["close"], timeperiod=n,
                                                    nbdevup=dev, nbdevdn=dev)
    if array:
        return upperband, middleband, lowerband
    return round(upperband[-1],3),round(middleband[-1],3), round(lowerband[-1],3)

# ---------------------------------------------------------
def boll2(data, n, devup, devdn, array=False):
    """布林通道, 上下轨可以不对称"""
    arrDict = transferData(data)  # 数据转化
    upperband, middleband, lowerband = talib.BBANDS(arrDict["close"], timeperiod=n,
                                                    nbdevup=devup, nbdevdn=devdn)
    if array:
        return upperband, middleband, lowerband
    return upperband[-1], middleband[-1], lowerband[-1]

# ---------------------------------------------------------
def donchian(data, n, array=False):
    """唐奇安通道"""
    arrDict = transferData(data)  # 数据转化
    up = talib.MAX(arrDict["high"], n)
    down = talib.MIN(arrDict["low"], n)
    if array:
        return up, down
    return up[-1], down[-1]

# ---------------------------------------------------------
def dx(data, n, array=False):
    """动向指标或趋向指标"""
    arrDict = transferData(data)  # 数据转化
    result = talib.DX(arrDict["high"], arrDict["low"], arrDict["close"], timeperiod=n)
    if array:
        return result
    return result[-1]

# ---------------------------------------------------------
def adx(data, n, array=False):
    """平均趋向指数"""
    arrDict = transferData(data)  # 数据转化
    result = talib.ADX(arrDict["high"], arrDict["low"], arrDict["close"], timeperiod=n)
    if array:
        return result
    return result[-1]

# ---------------------------------------------------------
def adxr(data, n, array=False):
    """平均趋向指数的趋向指数"""
    arrDict = transferData(data)  # 数据转化
    result = talib.ADXR(arrDict["high"], arrDict["low"], arrDict["close"], timeperiod=n)
    if array:
        return result
    return result[-1]

# ---------------------------------------------------------
def aroon(data, n, array=False):
    """阿隆指标"""
    arrDict = transferData(data)  # 数据转化
    aroondown, aroonup = talib.AROON(arrDict["high"], arrDict["low"], timeperiod=n)
    if array:
        return aroondown, aroonup
    return aroondown[-1], aroonup[-1]

# ---------------------------------------------------------
def aroonosc(data, n, array=False):
    """阿隆振荡指标"""
    arrDict = transferData(data)  # 数据转化
    result = talib.AROONOSC(arrDict["high"], arrDict["low"], timeperiod=n)
    if array:
        return result
    return result[-1]

# ---------------------------------------------------------
def cci(data, n, array=False):
    """顺势指标"""
    arrDict = transferData(data)  # 数据转化
    result = talib.CCI(arrDict["high"], arrDict["low"], arrDict["close"], timeperiod=n)
    if array:
        return result
    return result[-1]

# ---------------------------------------------------------
def cmo(data, n, array=False):
    """钱德动量摆动指标"""
    arrDict = transferData(data)  # 数据转化
    result = talib.CMO(arrDict["close"], timeperiod=n)
    if array:
        return result
    return result[-1]

# ---------------------------------------------------------
def macd(data, n1, n2, n3, array=False):
    """钱德动量摆动指标"""
    arrDict = transferData(data)  # 数据转化
    macdvalue, macdsignal, macdhist = talib.MACD(arrDict["close"], fastperiod=n1,
                                            slowperiod=n2, signalperiod=n3)
    if array:
        return macdvalue, macdsignal, macdhist
    return macdvalue[-1], macdsignal[-1], macdhist[-1]

# ---------------------------------------------------------
def mfi(data, n, array=False):
    """顺势指标"""
    arrDict = transferData(data)  # 数据转化
    result = talib.MFI(arrDict["high"], arrDict["low"], arrDict["close"], arrDict["volume"], timeperiod=n)
    if array:
        return result
    return result[-1]

# ---------------------------------------------------------
def mom(data, n, array=False):
    """动量指标"""
    arrDict = transferData(data)  # 数据转化
    result = talib.MOM(arrDict["close"], timeperiod=n)
    if array:
        return result
    return result[-1]

# ---------------------------------------------------------
def rsi(data, n, array=False):
    """相对强弱指数"""
    arrDict = transferData(data)  # 数据转化
    result = talib.RSI(arrDict["close"], timeperiod=n)
    if array:
        return result
    return result[-1]

# ---------------------------------------------------------
def kdj(data, n1, n2, n3, array=False):
    """
    KDJ指标, 其中n1为计算RSV的参数,
    n2为计算k的参数, n3为计算d的参数,一般为(9,3,3)
    """
    arrDict = transferData(data)  # 数据转化
    k, d = talib.STOCH(arrDict["high"], arrDict["low"], arrDict["close"],
                       fastk_period=n1, slowk_period=n2, slowd_period=n3)
    j = 3 * k - 2 * d
    if array:
        return k, d, j
    return k[-1], d[-1], j[-1]

# ---------------------------------------------------------
def willr(data, n, array=False):
    """威廉指标"""
    arrDict = transferData(data)  # 数据转化
    result = talib.WILLR(arrDict["high"], arrDict["low"], arrDict["close"], timeperiod=n)
    if array:
        return result
    return result[-1]

# ---------------------------------------------------------
def ad(data, array=False):
    """Chaikin A/D Line 累积/派发线"""
    arrDict = transferData(data)  # 数据转化
    result = talib.AD(arrDict["high"], arrDict["low"], arrDict["close"], arrDict["volume"])
    if array:
        return result
    return result[-1]

# ---------------------------------------------------------
def adosc(data, n1, n2, array=False):
    """Chaikin A/D Oscillator Chaikin震荡指标"""
    arrDict = transferData(data)  # 数据转化
    result = talib.ADOSC(arrDict["high"], arrDict["low"], arrDict["close"],
                         arrDict["volume"], fastperiod=n1, slowperiod=n2)
    if array:
        return result
    return result[-1]

# ---------------------------------------------------------
def obv(data, array=False):
    """On Balance Volume 能量潮"""
    arrDict = transferData(data)  # 数据转化
    result = talib.OBV(arrDict["close"], arrDict["volume"])
    if array:
        return result
    return result[-1]

# ---------------------------------------------------------
def atr(data, n, array=False):
    """真实波动幅度均值"""
    arrDict = transferData(data)  # 数据转化
    result = talib.ATR(arrDict["high"], arrDict["low"], arrDict["close"], timeperiod=n)
    if array:
        return result
    return result[-1]

# ---------------------------------------------------------
def trange(data, array=False):
    """真正的范围"""
    arrDict = transferData(data)  # 数据转化
    result = talib.TRANGE(arrDict["high"], arrDict["low"], arrDict["close"])
    if array:
        return result
    return result[-1]

# ---------------------------------------------------------
def avgprice(data, array=False):
    """平均价格函数"""
    arrDict = transferData(data)  # 数据转化
    result = talib.AVGPRICE(arrDict["open"], arrDict["high"], arrDict["low"], arrDict["close"])
    if array:
        return result
    return result[-1]

# ---------------------------------------------------------
def medprice(data, array=False):
    """中位数价格"""
    arrDict = transferData(data)  # 数据转化
    result = talib.MEDPRICE(arrDict["high"], arrDict["low"])
    if array:
        return result
    return result[-1]

# ---------------------------------------------------------
def typprice(data, array=False):
    """代表性价格"""
    arrDict = transferData(data)  # 数据转化
    result = talib.TYPPRICE(arrDict["high"], arrDict["low"], arrDict["close"])
    if array:
        return result
    return result[-1]

# ---------------------------------------------------------
def beta(data, n, array=False):
    """贝塔系数"""
    arrDict = transferData(data)  # 数据转化
    result = talib.BETA(arrDict["high"], arrDict["low"], timeperiod=n)
    if array:
        return result
    return result[-1]

# ---------------------------------------------------------
def linearreg(data, n, array=False):
    """线性回归估计值y"""
    arrDict = transferData(data)  # 数据转化
    result = talib.LINEARREG(arrDict["close"], timeperiod=n)
    if array:
        return result
    return result[-1]

# ---------------------------------------------------------
def linearreg_angle(data, n, array=False):
    """线性回归的角度, 单位为°"""
    arrDict = transferData(data)  # 数据转化
    result = talib.LINEARREG_ANGLE(arrDict["close"], timeperiod=n)
    if array:
        return result
    return result[-1]

# ---------------------------------------------------------
def linearreg_slope(data, n, array=False):
    """线性回归的斜率"""
    arrDict = transferData(data)  # 数据转化
    result = talib.LINEARREG_SLOPE(arrDict["close"], timeperiod=n)
    if array:
        return result
    return result[-1]

# ---------------------------------------------------------
def linearreg_intercept(data, n, array=False):
    """线性回归的截距"""
    arrDict = transferData(data)  # 数据转化
    result = talib.LINEARREG_INTERCEPT(arrDict["close"], timeperiod=n)
    if array:
        return result
    return result[-1]

# ---------------------------------------------------------
def std(data, n, array=False):
    """样本标准差"""
    arrDict = transferData(data)  # 数据转化
    result = talib.STDDEV(arrDict["close"], timeperiod=n, nbdev=1) # 默认是样本标准差
    if array:
        return result
    return result[-1]

# ---------------------------------------------------------
def var(data, n, array=False):
    """样本方差"""
    arrDict = transferData(data)  # 数据转化
    result = talib.VAR(arrDict["close"], timeperiod=n, nbdev=1)  # 默认是样本标准差
    if array:
        return result
    return result[-1]


    # # ----------------------------------------------------------------------
    # def boll(self, n, dev, array=False):
    #     """布林通道"""
    #     mid = self.sma(n, array)
    #     std = self.std(n, array)
    #
    #     up = mid + std * dev
    #     down = mid - std * dev
    #
    #     return up, down, mid
